#include "ompl_extension/TrajectoryEvaluator.h"


namespace ompl
{
  namespace base
  {

    SmoothnessEvaluator::SmoothnessEvaluator(ompl::base::SpaceInformation* si,
                                             planning_environment::CollisionModelsInterface* cmi,
                                             const std::string& group_name,
                                             double discretization,
                                             double smoothness_cost_velocity,
                                             double smoothness_cost_acceleration,
                                             double smoothness_cost_jerk) : TrajectoryEvaluator()
  {
    discretization_ = discretization;
    smoothness_cost_velocity_ = smoothness_cost_velocity;
    smoothness_cost_acceleration_ = smoothness_cost_acceleration;
    smoothness_cost_jerk_ = smoothness_cost_jerk;
    si_ = si;
    cmi_ = cmi;

    joint_group_ = cmi->getKinematicModel()->getModelGroup(group_name);

    if(!ompl_ros_interface_helper::jointGroupToOmplStateSpacePtr(joint_group_, ompl_state_to_kinematic_state_mapping_, kinematic_state_to_ompl_state_mapping_))
    {
      ROS_ERROR("Could not set up the ompl state space");
      return;
    }
  }

    double SmoothnessEvaluator::getCost(const std::vector<State*>& trajectory) const
    {
      unsigned int n_points = trajectory.size();
      unsigned int dim = si_->getStateDimension();

      // compute the quality matrix
      int num_vars_all = n_points;
      Eigen::MatrixXd diff_matrix = Eigen::MatrixXd::Zero(num_vars_all, num_vars_all);
      Eigen::MatrixXd quad_cost_matrix = Eigen::MatrixXd::Zero(num_vars_all, num_vars_all);

      double multiplier = 1.0;

      multiplier *= discretization_;
      diff_matrix = getDiffMatrix(num_vars_all, &DIFF_RULES[0][0]);
      quad_cost_matrix += smoothness_cost_velocity_ * multiplier * diff_matrix.transpose() * diff_matrix;

      multiplier *= discretization_;
      diff_matrix = getDiffMatrix(num_vars_all, &DIFF_RULES[1][0]);
      quad_cost_matrix += smoothness_cost_acceleration_ * multiplier * diff_matrix.transpose() * diff_matrix;

      multiplier *= discretization_;
      diff_matrix = getDiffMatrix(num_vars_all, &DIFF_RULES[2][0]);
      quad_cost_matrix += smoothness_cost_jerk_ * multiplier * diff_matrix.transpose() * diff_matrix;

      planning_models::KinematicState* kinematic_state = cmi_->getPlanningSceneState();
      if(!kinematic_state)
      {
        ROS_ERROR_STREAM("Planning scene hasn't been set");
        return 0;
      }

      std::vector<std::vector<double> > joint_state_values(n_points);
      planning_models::KinematicState::JointStateGroup* joint_state_group = kinematic_state->getJointStateGroup(joint_group_->getName());

      for(unsigned int i = 0; i < n_points; ++i)
      {
        omplStateToKinematicStateGroup(trajectory[i], ompl_state_to_kinematic_state_mapping_, joint_state_group);
        joint_state_group->getKinematicStateValues(joint_state_values[i]);
      }


      double smoothness_cost = 0.0;

      Eigen::VectorXd v(n_points);

      for(unsigned int d = 0; d < dim; ++d)
      {
        for(unsigned int i = 0; i < n_points; ++i)
        {
          v(i) = joint_state_values[i][d];
        }

        smoothness_cost += v.dot(quad_cost_matrix * v);
      }

      return smoothness_cost;
    }

    Eigen::MatrixXd SmoothnessEvaluator::getDiffMatrix(int size, const double* diff_rule) const
    {
      Eigen::MatrixXd matrix = Eigen::MatrixXd::Zero(size, size);
      for(int i = 0; i < size; ++i)
      {
        for(int j = -DIFF_RULE_LENGTH/2; j <= DIFF_RULE_LENGTH/2; ++j)
        {
          int index = i + j;
          if(index < 0)
            continue;
          if(index >= size)
            continue;
          matrix(i, index) = diff_rule[j + DIFF_RULE_LENGTH/2];
        }
      }
      return matrix;
    }


    DynamicsEvaluator::DynamicsEvaluator(ompl::base::SpaceInformation* si,
                                         planning_environment::CollisionModelsInterface* cmi,
                                         const std::string& group_name,
                                         double discretization) : TrajectoryEvaluator()
    {
      discretization_ = discretization;
      si_ = si;
      cmi_ = cmi;

      joint_group_ = cmi->getKinematicModel()->getModelGroup(group_name);

      if(!ompl_ros_interface_helper::jointGroupToOmplStateSpacePtr(joint_group_, ompl_state_to_kinematic_state_mapping_, kinematic_state_to_ompl_state_mapping_))
      {
        ROS_ERROR("Could not set up the ompl state space");
        return;
      }


      if(!kdl_parser::treeFromUrdfModel(*(cmi->getParsedDescription()), kdl_tree_))
      {
        ROS_ERROR("Cound not construct KDL tree from URDF.");
        return;
      }

      if(!calculateKDLChain())
        return;

      // compute mapping from kdl joint index and the kinematic joint index
      const std::vector<const JointModel*>& all_joints = joint_group_->getJointModels();
      std::vector<std::string> joint_names(all_joints.size());
      for(unsigned int i = 0; i < all_joints.size(); ++i)
        joint_names[i] = all_joints[i]->getName();

      for(unsigned int i = 0; i < kdl_chain_.segments.size(); ++i)
      {
        if(kdl_chain_.segments[i].getJoint().getType() != KDL::Joint::None)
        {
          std::string kdl_joint_name = kdl_chain_.segments[i].getJoint().getName();
          kdl_joint_id_to_joint_name_mapping_.push_back(kdl_joint_name);

          if(std::find(joint_names.begin(), joint_names.end(), kdl_joint_name) != joint_names.end())
          {
            unsigned int id = std::find(joint_names.begin(), joint_names.end(), kdl_joint_name) - joint_names.begin();
            kdl_joint_id_to_joint_id_mapping_.insert(std::make_pair(i, id));
          }
          else
          {
            kdl_joint_id_to_joint_id_mapping_.insert(std::make_pair(i, -1));
          }
        }
      }

      KDL::Vector gravity(0, 0, -9.8);
      inverse_dynamics_solver_.reset(new KDL::ChainIdSolver_RNE(kdl_chain_, gravity));
    }

    bool DynamicsEvaluator::calculateKDLChain()
    {
      const std::vector<const JointModel*>& all_joints = joint_group_->getJointModels();

       std::vector<const JointModel*> end_effectors; // the end_effectors
       std::vector<const JointModel*> end_effector_roots; // the roots corresponding to each end-effector

       // compute all end-effectors, though in fact later we only allow one end-effector
       std::map<const JointModel*, bool> is_end_effectors;
       for(unsigned int i = 0; i < all_joints.size(); ++i)
       {
         const JointModel* joint = all_joints[i];
         is_end_effectors.insert(std::make_pair(joint, true));
       }

       for(unsigned int i = 0; i < all_joints.size(); ++i)
       {
         const JointModel* joint = all_joints[i];
         if(is_end_effectors[joint])
         {
           while(joint->getParentLinkModel())
           {
             joint = joint->getParentLinkModel()->getParentJointModel();
             if(joint_group_->hasJointModel(joint->getName()))
               is_end_effectors[joint] = false;
           }
         }
       }

       for(std::map<const JointModel*, bool>::iterator it = is_end_effectors.begin();
           it != is_end_effectors.end();
           ++it)
       {
         if(it->second)  end_effectors.push_back(it->first);
       }


       // get the root of the end effector
       end_effector_roots.resize(end_effectors.size());
       for(unsigned int i = 0; i < end_effectors.size(); ++i)
       {
         const JointModel* joint = end_effectors[i];

         end_effector_roots[i] = NULL;
         while(joint->getParentLinkModel())
         {
           joint = joint->getParentLinkModel()->getParentJointModel();
           if(joint_group_->hasJointModel(joint->getName()))
             end_effector_roots[i] = joint;
         }
       }

       for(unsigned int i = 0; i < end_effectors.size(); ++i)
         end_effector_roots[i] = end_effector_roots[i]->getParentLinkModel()->getParentJointModel();


       if(end_effectors.size() > 1)
       {
         ROS_ERROR("One group can only have one end-effector.");
         return false;
       }


       kdl_tree_.getChain(end_effector_roots[0]->getChildLinkModel()->getName(),
                          end_effectors[0]->getChildLinkModel()->getName(),
                          kdl_chain_);

       return true;
    }

    double DynamicsEvaluator::getCost(const std::vector<State*>& trajectory) const
    {
      unsigned int n_points = trajectory.size();

      if(n_points < DIFF_RULE_LENGTH)
      {
        std::cout << "the trajectory is too short: " << n_points << std::endl;
        return 0.0;
      }

      planning_models::KinematicState* kinematic_state = cmi_->getPlanningSceneState();
      if(!kinematic_state)
      {
        ROS_ERROR_STREAM("Planning scene hasn't been set");
        return 0;
      }

      std::vector<std::vector<double> > joint_state_values(n_points);
      planning_models::KinematicState::JointStateGroup* joint_state_group = kinematic_state->getJointStateGroup(joint_group_->getName());


      for(unsigned int i = 0; i < n_points; ++i)
      {
        omplStateToKinematicStateGroup(trajectory[i], ompl_state_to_kinematic_state_mapping_, joint_state_group);
        joint_state_group->getKinematicStateValues(joint_state_values[i]);
      }

      unsigned int nsegments = kdl_chain_.getNrOfSegments();

      unsigned int njoints = kdl_chain_.getNrOfJoints();

      const std::vector<KDL::Wrench> fext(nsegments);
      for(unsigned int i = 0; i < nsegments; ++i)
        fext[i].Zero();

      std::vector<std::vector<double> > joint_torque_values(n_points);
      for(unsigned int i = 0; i < n_points; ++i)
        joint_torque_values[i].resize(njoints, 0);


      for(unsigned int i = DIFF_RULE_LENGTH/2; i < n_points - DIFF_RULE_LENGTH/2; ++i)
        getTorques(i, joint_state_values, fext, joint_torque_values[i]);


      // return the cost as the maximum torque over time
      double torque_cost = 0;

      for(unsigned int i = DIFF_RULE_LENGTH/2; i < n_points - DIFF_RULE_LENGTH/2; ++i)
      {
        double cost = 0;
        for(unsigned int j = 0; j < njoints; ++j)
          cost += fabs(joint_torque_values[i][j]);

        if(cost > torque_cost) torque_cost = cost;
      }

      return torque_cost;
    }

    void DynamicsEvaluator::getTorques(int traj_point,
                                       const std::vector<std::vector<double> >& joint_state_values,
                                       const std::vector<KDL::Wrench>& fext,
                                       std::vector<double>& torques) const
    {
      unsigned int njoints = kdl_chain_.getNrOfJoints();

      KDL::JntArray kdl_jnt_positions, kdl_jnt_velocities, kdl_jnt_accelerations;

      getJointPositions(traj_point, joint_state_values, kdl_jnt_positions);

      getJointVelocities(traj_point, joint_state_values, kdl_jnt_velocities);

      getJointAccelerations(traj_point, joint_state_values, kdl_jnt_accelerations);

      KDL::JntArray kdl_jnt_torques(njoints);

      inverse_dynamics_solver_->CartToJnt(kdl_jnt_positions,
                                          kdl_jnt_velocities,
                                          kdl_jnt_accelerations,
                                          fext,
                                          kdl_jnt_torques);

      // not transform from segment id order to joint id order :(
      for(unsigned int i = 0; i < njoints; ++i)
        torques[i] = kdl_jnt_torques(i);
    }

    void DynamicsEvaluator::getJointPositions(int traj_point,
                                              const std::vector<std::vector<double> >& joint_state_values,
                                              KDL::JntArray& kdl_jnt_positions) const
    {
      unsigned int njoints = kdl_chain_.getNrOfJoints();

      kdl_jnt_positions.resize(njoints);

      for(unsigned int i = 0; i < njoints; ++i)
      {
        kdl_jnt_positions(i) = 0;
        int id = (kdl_joint_id_to_joint_id_mapping_.find(i))->second;
        if(id != -1)
          kdl_jnt_positions(i) = joint_state_values[traj_point][id];
        else
          kdl_jnt_positions(i) = cmi_->getPlanningSceneState()->getJointState(kdl_joint_id_to_joint_name_mapping_[i])->getJointStateValues()[0];
      }
    }


    void DynamicsEvaluator::getJointVelocities(int traj_point,
                                               const std::vector<std::vector<double> >& joint_state_values,
                                               KDL::JntArray& kdl_jnt_velocities) const
    {
      unsigned int njoints = kdl_chain_.getNrOfJoints();

      kdl_jnt_velocities.resize(njoints);

      double invTime = 1.0 / discretization_;

      for(unsigned int i = 0; i < njoints; ++i)
      {
        kdl_jnt_velocities(i) = 0;
        int id = (kdl_joint_id_to_joint_id_mapping_.find(i))->second;
        if(id != -1)
        {
          for(int k = -DIFF_RULE_LENGTH/2; k <= DIFF_RULE_LENGTH/2; ++k)
            kdl_jnt_velocities(i) += (invTime * DIFF_RULES[0][k + DIFF_RULE_LENGTH/2]) * joint_state_values[traj_point + k][id];
        }
      }
    }

    void DynamicsEvaluator::getJointAccelerations(int traj_point,
                                                  const std::vector<std::vector<double> >& joint_state_values,
                                                  KDL::JntArray& kdl_jnt_accelerations) const
    {
      unsigned int njoints = kdl_chain_.getNrOfJoints();

      kdl_jnt_accelerations.resize(njoints);

      double invTime = 1.0 / (discretization_ * discretization_);

      for(unsigned int i = 0; i < njoints; ++i)
      {
        kdl_jnt_accelerations(i) = 0;
        int id = (kdl_joint_id_to_joint_id_mapping_.find(i))->second;
        if(id != -1)
        {
          for(int k = -DIFF_RULE_LENGTH/2; k <= DIFF_RULE_LENGTH/2; ++k)
            kdl_jnt_accelerations(i) += (invTime * DIFF_RULES[1][k + DIFF_RULE_LENGTH/2]) * joint_state_values[traj_point + k][id];
        }
      }
    }



  }
}
